# 仿真器入口 用于训练强化学习 跟踪算法用的简化版本仿真器
# 差速控制器训练
# 航控需要通过 速度偏差和航向偏差 等指标查看控制
from dynamics_simulation import motionSim
import numpy as np
import random
import socket
import json
import traceback
from algo_struct.algo import rad_limit
from math import pow

# 解析控制参数
def decode_data(data):
    # load json
    recv_obser = json.loads(data)
    # print(recv_obser)
    rspl = recv_obser['rspl']
    rspr = recv_obser['rspr']

    return rspl, rspr

# 打包船的状态返回
def encode_data(trans):
    # 按照标准接口编写
    data = json.dumps(trans, sort_keys=True, indent=4, separators=(',', ':'))
    return data.encode('utf-8')

# 计算奖励
def get_reward(observation: dict):
    score =100 - ( pow(180.0/np.pi * rad_limit((observation['exp_psi'] - observation['psi'])*np.pi/180.0) ,2)
                + 100 * pow(observation['exp_speed'] - observation['v_x'], 2)
                + observation['v_r']**2
                + 100 * observation['acc_x']**2
                + 100 * observation['acc_y']**2
                + 1e-7 * observation['rspl']**2
                + 1e-7 * observation['rspr']**2)

    print(f'score is {score}')
    return score

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--port', type=int, default=11023)
args = parser.parse_args()

# test motion sim
if __name__ == '__main__':
    boat_data = dict()
    # 与航行控制算法连接TCP
    LOCAL_HOST = "127.0.0.1"
    LOCAL_PORT = args.port
    
    socket_server = socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM)
    socket_server.bind((LOCAL_HOST, LOCAL_PORT))
    socket_server.listen()
    conn, address = socket_server.accept()
    print(f"收到了客户端的连接，客户端信息是 {address}")

    all_sim_time = 200
    sim_time = 0
    delta_h = 0.2
    
    # boat运动仿真类输入参数
    my_mmg = motionSim.BoatMotionSim(delta_h=delta_h, name='usv')
    
    delta = 0
    try:
        # 给出期望值 期望航向-180～180 期望航速+5～-2
        # exp_speed = random.randint(-2, 3)
        exp_speed = random.choice([-2, -1, 0, 1, 2])# 指定list中随机选择
        exp_angle = random.randint(-90, 90)
        x = np.array([0.0, 0.0, 0.0])# 初始无人艇位置
        u = np.array([0.0, 0.0, 0.0])# 初始无人艇速度
        V = 0
        last_u = 0.0

        delta = 0
        rsp = 0
        while True:
            sim_time += delta_h
            print(f'期望speed和angle {exp_speed}, {exp_angle}')

            # 接收转速和舵角
            info, address = conn.recvfrom(1024)
            rspl, rspr = decode_data(info)
            print(f'驱动力和舵角 {rspl}, {rspr}')

            def rsp_to_force(rsp):
                a = 2.30549953e-04
                b = 3.88991253e-01
                c = -2.87920929e+02
                # 2.30549953e-04  3.88991253e-01 -2.87920929e+02
                return a*rsp*abs(rsp) + b*rsp + c

            # rsp to force
            f_left = rsp_to_force(rspl)
            f_right = rsp_to_force(rspr)
            # 计算更新 位置 速度
            x, u, nu, acc = my_mmg.run_MMG_double_thrust(x=x, u=u, f_left=f_left, f_right=f_right)
            
            x[2] = rad_limit(x[2]) # 限制到-pi到pi 
            
            V = u[0]/np.abs(u[0]) * np.sqrt(u[0]**2 + u[1]**2 + 1e-6)
            
            # 为强化学习用
            observation=dict()
            observation['exp_psi'] = exp_angle # 期望航向
            observation['exp_speed'] = exp_speed # 期望速度 m/s
            observation['rspl'] = rspl # 转速左 rpm
            observation['rspr'] = rspr # 转速右 rpm
            observation['psi'] = x[2]* 180.0/np.pi # 位置psi
            observation['v_r'] = u[2]* 180.0/np.pi # z轴角速度 degree/s
            observation['v_x'] = u[0] # 船体坐标系x速度
            observation['v_y'] = u[1] # 船体坐标系y速度
            observation['V'] = V # 合速度
            observation['acc_x'] = acc[0] # 加速度
            observation['acc_y'] = acc[1]
            observation['acc_r'] = acc[2]
            observation['delta_psi'] = 180.0/np.pi * rad_limit(observation['exp_psi']*np.pi/180.0 - observation['psi']*np.pi/180.0)# 偏差角度
            observation['delta_speed'] = observation['exp_speed'] - observation['V']# 偏差速度

            boat_data['reward'] = get_reward(observation)
            print(f'observation is {observation}')
            
            if sim_time > all_sim_time:
                # 给出期望值 期望航向-180～180 期望航速-1～1
                exp_speed = random.choice([-2, -1, 0, 1, 2])# 指定list中随机选择
                exp_angle = random.randint(-90, 90)
                x = np.array([0.0, 0.0, 0.0])# 初始无人艇位置
                u = np.array([0.0, 0.0, 0.0])# 初始无人艇速度
                rspl = 0 # 转速
                rspr = 0 # 舵角
                sim_time = 0
                done = 1
            else:
                done = 0

            boat_data['observation'] = observation
            boat_data['terminated'] = done
            boat_data['truncated'] = done
            data = encode_data(trans=boat_data)
            # 发送数据
            conn.send(data)

    except Exception:
        # 关闭连接
        conn.close()
        socket_server.close()

        traceback.print_exc()
